#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import sys
import datetime
import os


def SearchFileInMongo(collection, filepath):
    result = collection.find_one({"path": filepath})
    if result != None:
        return True
    else:
        return False


def SearchFilesInMongo(collection, filepath):
    result = collection.find_one({"paths": filepath})
    if result != None:
        return True
    else:
        return False


if __name__ == "__main__":
    infile = open(sys.argv[1], "r")

    ftp_head = "ftp://ftp.ncbi.nlm.nih.gov/pub/pmc/"

    DB_NAME = "pmc"
    COLLECTION_NAME = "filelist"

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    file_time = infile.readline().rstrip()
    split_file_date = file_time.split(" ")[0].split("-")
    split_file_time = file_time.split(" ")[1].split(":")
    year = int(split_file_date[0])
    month = int(split_file_date[1])
    day = int(split_file_date[2])
    hour = int(split_file_time[0])
    minute = int(split_file_time[1])
    second = int(split_file_time[2])
    in_time = datetime.datetime(year, month, day, hour, minute, second)

    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        path = splitline[0]
        type = ""
        exactpath = ""
        if path.startswith("oa_package"):
            exactpath = path[11:]
            type = "tarball"
        elif path.startswith("oa_pdf"):
            exactpath = path[7:]
            type = "pdf"

        if not SearchFileInMongo(collection, exactpath):
            if type == "pdf":
                dataset = {}
                name = splitline[1]
                if len(splitline) >= 3:
                    if splitline[2] != "":
                        pmcid = int(splitline[2][3:])
                        dataset["pmcid"] = pmcid
                if len(splitline) >= 4:
                    if splitline[3] != "":
                        pmid = int(splitline[3].split(":")[1])
                        dataset["pmid"] = pmid
                dataset["name"] = name
                if exactpath != "":
                    dataset["path"] = exactpath
                if type != "":
                    dataset["type"] = type
                dataset["time_in"] = in_time
                collection.insert_one(dataset)
                print ftp_head + path
            elif type == "tarball":
                result = collection.find_one({"paths": exactpath})
                if result != None:
                    id = result["_id"]
                    old_path = "pmc_pdf/" + result["path"]
                    new_path = "pmc_pdf/" + exactpath
                    os.rename(old_path, new_path)
                    collection.update_one({"_id": id}, {"$set": {"path": exactpath}})
                else:
                    dataset = {}
                    name = splitline[1]
                    if len(splitline) >= 3:
                        if splitline[2] != "":
                            pmcid = int(splitline[2][3:])
                            dataset["pmcid"] = pmcid
                            id_path = exactpath[:6] + "PMC" + str(pmcid) + ".tar.gz"
                            if exactpath != id_path:
                                dataset["paths"] = [exactpath, id_path]
                            else:
                                dataset["paths"] = [exactpath]
                    if len(splitline) >= 4:
                        if splitline[3] != "":
                            pmid = int(splitline[3].split(":")[1])
                            dataset["pmid"] = pmid
                    dataset["name"] = name
                    if exactpath != "":
                        dataset["path"] = exactpath
                    if type != "":
                        dataset["type"] = type
                    dataset["time_in"] = in_time
                    collection.insert_one(dataset)
                    print ftp_head + path

    infile.close()
    connection.close()
